# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================

"""A set of functions that are used for visualization.

These functions often receive an image, perform some visualization on the image.
The functions do not return a value, instead they modify the image itself.

"""
import os
import collections
import functools
import matplotlib.pyplot as plt
import numpy as np
import PIL.Image as Image
import PIL.ImageColor as ImageColor
import PIL.ImageDraw as ImageDraw
import PIL.ImageFont as ImageFont
import six

from .common import get_default_font

_TITLE_LEFT_MARGIN = 10
_TITLE_TOP_MARGIN = 10
STANDARD_COLORS = [
    'AliceBlue', 'Chartreuse', 'Aqua', 'Aquamarine', 'Azure', 'Beige', 'Bisque',
    'BlanchedAlmond', 'BlueViolet', 'BurlyWood', 'CadetBlue', 'AntiqueWhite',
    'Chocolate', 'Coral', 'CornflowerBlue', 'Cornsilk', 'Crimson', 'Cyan',
    'DarkCyan', 'DarkGoldenRod', 'DarkGrey', 'DarkKhaki', 'DarkOrange',
    'DarkOrchid', 'DarkSalmon', 'DarkSeaGreen', 'DarkTurquoise', 'DarkViolet',
    'DeepPink', 'DeepSkyBlue', 'DodgerBlue', 'FireBrick', 'FloralWhite',
    'ForestGreen', 'Fuchsia', 'Gainsboro', 'GhostWhite', 'Gold', 'GoldenRod',
    'Salmon', 'Tan', 'HoneyDew', 'HotPink', 'IndianRed', 'Ivory', 'Khaki',
    'Lavender', 'LavenderBlush', 'LawnGreen', 'LemonChiffon', 'LightBlue',
    'LightCoral', 'LightCyan', 'LightGoldenRodYellow', 'LightGray', 'LightGrey',
    'LightGreen', 'LightPink', 'LightSalmon', 'LightSeaGreen', 'LightSkyBlue',
    'LightSlateGray', 'LightSlateGrey', 'LightSteelBlue', 'LightYellow', 'Lime',
    'LimeGreen', 'Linen', 'Magenta', 'MediumAquaMarine', 'MediumOrchid',
    'MediumPurple', 'MediumSeaGreen', 'MediumSlateBlue', 'MediumSpringGreen',
    'MediumTurquoise', 'MediumVioletRed', 'MintCream', 'MistyRose', 'Moccasin',
    'NavajoWhite', 'OldLace', 'Olive', 'OliveDrab', 'Orange', 'OrangeRed',
    'Orchid', 'PaleGoldenRod', 'PaleGreen', 'PaleTurquoise', 'PaleVioletRed',
    'PapayaWhip', 'PeachPuff', 'Peru', 'Pink', 'Plum', 'PowderBlue', 'Purple',
    'Red', 'RosyBrown', 'RoyalBlue', 'SaddleBrown', 'Green', 'SandyBrown',
    'SeaGreen', 'SeaShell', 'Sienna', 'Silver', 'SkyBlue', 'SlateBlue',
    'SlateGray', 'SlateGrey', 'Snow', 'SpringGreen', 'SteelBlue', 'GreenYellow',
    'Teal', 'Thistle', 'Tomato', 'Turquoise', 'Violet', 'Wheat', 'White',
    'WhiteSmoke', 'Yellow', 'YellowGreen'
]


def save_image_array_as_png(image, output_path):
    """Saves an image (represented as a numpy array) to PNG.

    Args:
      image: a numpy array with shape [height, width, 3].
      output_path: path to which image should be written.
    """
    image_pil = Image.fromarray(np.uint8(image)).convert('RGB')
    with open(output_path, 'wb') as fid:
        image_pil.save(fid, 'PNG')


def encode_image_array_as_png_str(image):
    """Encodes a numpy array into a PNG string.

    Args:
      image: a numpy array with shape [height, width, 3].

    Returns:
      PNG encoded image string.
    """
    image_pil = Image.fromarray(np.uint8(image))
    output = six.BytesIO()
    image_pil.save(output, format='PNG')
    png_string = output.getvalue()
    output.close()
    return png_string


def draw_bounding_box_on_image_array(image,
                                     ymin,
                                     xmin,
                                     ymax,
                                     xmax,
                                     color='red',
                                     thickness=4,
                                     display_str_list=(),
                                     use_normalized_coordinates=True):
    """Adds a bounding box to an image (numpy array).

    Bounding box coordinates can be specified in either absolute (pixel) or
    normalized coordinates by setting the use_normalized_coordinates argument.

    Args:
      image: a numpy array with shape [height, width, 3].
      ymin: ymin of bounding box.
      xmin: xmin of bounding box.
      ymax: ymax of bounding box.
      xmax: xmax of bounding box.
      color: color to draw bounding box. Default is red.
      thickness: line thickness. Default value is 4.
      display_str_list: list of strings to display in box
                        (each to be shown on its own line).
      use_normalized_coordinates: If True (default), treat coordinates
        ymin, xmin, ymax, xmax as relative to the image.  Otherwise treat
        coordinates as absolute.
    """
    image_pil = Image.fromarray(np.uint8(image)).convert('RGB')
    draw_bounding_box_on_image(image_pil, ymin, xmin, ymax, xmax, color,
                               thickness, display_str_list,
                               use_normalized_coordinates)
    np.copyto(image, np.array(image_pil))


def draw_bounding_box_on_image(image,
                               ymin,
                               xmin,
                               ymax,
                               xmax,
                               color='red',
                               thickness=4,
                               display_str_list=(),
                               use_normalized_coordinates=True):
    """Adds a bounding box to an image.

    Bounding box coordinates can be specified in either absolute (pixel) or
    normalized coordinates by setting the use_normalized_coordinates argument.

    Each string in display_str_list is displayed on a separate line above the
    bounding box in black text on a rectangle filled with the input 'color'.
    If the top of the bounding box extends to the edge of the image, the strings
    are displayed below the bounding box.

    Args:
      image: a PIL.Image object.
      ymin: ymin of bounding box.
      xmin: xmin of bounding box.
      ymax: ymax of bounding box.
      xmax: xmax of bounding box.
      color: color to draw bounding box. Default is red.
      thickness: line thickness. Default value is 4.
      display_str_list: list of strings to display in box
                        (each to be shown on its own line).
      use_normalized_coordinates: If True (default), treat coordinates
        ymin, xmin, ymax, xmax as relative to the image.  Otherwise treat
        coordinates as absolute.
    """
    draw = ImageDraw.Draw(image)
    im_width, im_height = image.size
    if use_normalized_coordinates:
        (left, right, top, bottom) = (xmin * im_width, xmax * im_width,
                                      ymin * im_height, ymax * im_height)
    else:
        (left, right, top, bottom) = (xmin, xmax, ymin, ymax)
    draw.line([(left, top), (left, bottom), (right, bottom),
               (right, top), (left, top)], width=thickness, fill=color)

    font = get_default_font(24)

    # If the total height of the display strings added to the top of the bounding
    # box exceeds the top of the image, stack the strings below the bounding box
    # instead of above.
    display_str_heights = [font.getsize(ds)[1] for ds in display_str_list]
    # Each display_str has a top and bottom margin of 0.05x.
    total_display_str_height = (1 + 2 * 0.05) * sum(display_str_heights)

    if top > total_display_str_height:
        text_bottom = top
    else:
        text_bottom = bottom + total_display_str_height
    # Reverse list and print from bottom to top.
    for display_str in display_str_list[::-1]:
        text_width, text_height = font.getsize(display_str)
        margin = np.ceil(0.05 * text_height)
        draw.rectangle(
            [(left, text_bottom - text_height - 2 * margin), (left + text_width,
                                                              text_bottom)],
            fill=color)
        draw.text(
            (left + margin, text_bottom - text_height - margin),
            display_str,
            fill='black',
            font=font)
        text_bottom -= text_height - 2 * margin


def draw_bounding_boxes_on_image_array(image,
                                       boxes,
                                       color='red',
                                       thickness=4,
                                       display_str_list_list=()):
    """Draws bounding boxes on image (numpy array).

    Args:
      image: a numpy array object.
      boxes: a 2 dimensional numpy array of [N, 4]: (ymin, xmin, ymax, xmax).
             The coordinates are in normalized format between [0, 1].
      color: color to draw bounding box. Default is red.
      thickness: line thickness. Default value is 4.
      display_str_list_list: list of list of strings.
                             a list of strings for each bounding box.
                             The reason to pass a list of strings for a
                             bounding box is that it might contain
                             multiple labels.

    Raises:
      ValueError: if boxes is not a [N, 4] array
    """
    image_pil = Image.fromarray(image)
    draw_bounding_boxes_on_image(image_pil, boxes, color, thickness,
                                 display_str_list_list)
    np.copyto(image, np.array(image_pil))


def draw_bounding_boxes_on_image(image,
                                 boxes,
                                 color='red',
                                 thickness=4,
                                 display_str_list_list=()):
    """Draws bounding boxes on image.

    Args:
      image: a PIL.Image object.
      boxes: a 2 dimensional numpy array of [N, 4]: (ymin, xmin, ymax, xmax).
             The coordinates are in normalized format between [0, 1].
      color: color to draw bounding box. Default is red.
      thickness: line thickness. Default value is 4.
      display_str_list_list: list of list of strings.
                             a list of strings for each bounding box.
                             The reason to pass a list of strings for a
                             bounding box is that it might contain
                             multiple labels.

    Raises:
      ValueError: if boxes is not a [N, 4] array
    """
    boxes_shape = boxes.shape
    if not boxes_shape:
        return
    if len(boxes_shape) != 2 or boxes_shape[1] != 4:
        raise ValueError('Input must be of size [N, 4]')
    for i in range(boxes_shape[0]):
        display_str_list = ()
        if display_str_list_list:
            display_str_list = display_str_list_list[i]
        draw_bounding_box_on_image(image, boxes[i, 0], boxes[i, 1], boxes[i, 2],
                                   boxes[i, 3], color, thickness, display_str_list)


def draw_keypoints_on_image_array(image,
                                  keypoints,
                                  color='red',
                                  radius=2,
                                  use_normalized_coordinates=True):
    """Draws keypoints on an image (numpy array).

    Args:
      image: a numpy array with shape [height, width, 3].
      keypoints: a numpy array with shape [num_keypoints, 2].
      color: color to draw the keypoints with. Default is red.
      radius: keypoint radius. Default value is 2.
      use_normalized_coordinates: if True (default), treat keypoint values as
        relative to the image.  Otherwise treat them as absolute.
    """
    image_pil = Image.fromarray(np.uint8(image)).convert('RGB')
    draw_keypoints_on_image(image_pil, keypoints, color, radius,
                            use_normalized_coordinates)
    np.copyto(image, np.array(image_pil))


def draw_keypoints_on_image(image,
                            keypoints,
                            color='red',
                            radius=2,
                            use_normalized_coordinates=True):
    """Draws keypoints on an image.

    Args:
      image: a PIL.Image object.
      keypoints: a numpy array with shape [num_keypoints, 2].
      color: color to draw the keypoints with. Default is red.
      radius: keypoint radius. Default value is 2.
      use_normalized_coordinates: if True (default), treat keypoint values as
        relative to the image.  Otherwise treat them as absolute.
    """
    draw = ImageDraw.Draw(image)
    im_width, im_height = image.size
    keypoints_x = [k[1] for k in keypoints]
    keypoints_y = [k[0] for k in keypoints]
    if use_normalized_coordinates:
        keypoints_x = tuple([im_width * x for x in keypoints_x])
        keypoints_y = tuple([im_height * y for y in keypoints_y])
    for keypoint_x, keypoint_y in zip(keypoints_x, keypoints_y):
        draw.ellipse([(keypoint_x - radius, keypoint_y - radius),
                      (keypoint_x + radius, keypoint_y + radius)],
                     outline=color, fill=color)


def draw_mask_on_image_array(image, mask, color='red', alpha=0.4):
    """Draws mask on an image.

    Args:
      image: uint8 numpy array with shape (img_height, img_height, 3)
      mask: a uint8 numpy array of shape (img_height, img_height) with
        values between either 0 or 1.
      color: color to draw the keypoints with. Default is red.
      alpha: transparency value between 0 and 1. (default: 0.4)

    Raises:
      ValueError: On incorrect data type for image or masks.
    """
    if image.dtype != np.uint8:
        raise ValueError('`image` not of type np.uint8')
    if mask.dtype != np.uint8:
        raise ValueError('`mask` not of type np.uint8')
    if np.any(np.logical_and(mask != 1, mask != 0)):
        raise ValueError('`mask` elements should be in [0, 1]')
    if image.shape[:2] != mask.shape:
        raise ValueError('The image has spatial dimensions %s but the mask has '
                         'dimensions %s' % (image.shape[:2], mask.shape))
    rgb = ImageColor.getrgb(color)
    pil_image = Image.fromarray(image)

    solid_color = np.expand_dims(
        np.ones_like(mask), axis=2) * np.reshape(list(rgb), [1, 1, 3])
    pil_solid_color = Image.fromarray(np.uint8(solid_color)).convert('RGBA')
    pil_mask = Image.fromarray(np.uint8(255.0 * alpha * mask)).convert('L')
    pil_image = Image.composite(pil_solid_color, pil_image, pil_mask)
    np.copyto(image, np.array(pil_image.convert('RGB')))


def visualize_boxes_and_labels_on_image_array(
    image,
    boxes,
    classes,
    scores,
    category_index,
    instance_masks=None,
    instance_boundaries=None,
    keypoints=None,
    use_normalized_coordinates=False,
    max_boxes_to_draw=None,
    min_score_thresh=0.,
    agnostic_mode=False,
    line_thickness=4,
    groundtruth_box_visualization_color='black',
    skip_scores=False,
    skip_labels=False):
    """Overlay labeled boxes on an image with formatted scores and label names.

    This function groups boxes that correspond to the same location
    and creates a display string for each detection and overlays these
    on the image. Note that this function modifies the image in place, and returns
    that same image.

    Args:
      image: uint8 numpy array with shape (img_height, img_width, 3)
      boxes: a numpy array of shape [N, 4]
      classes: a numpy array of shape [N]. Note that class indices are 1-based,
        and match the keys in the label map.
      scores: a numpy array of shape [N] or None.  If scores=None, then
        this function assumes that the boxes to be plotted are groundtruth
        boxes and plot all boxes as black with no classes or scores.
      category_index: a dict containing category dictionaries (each holding
        category index `id` and category name `name`) keyed by category indices.
      instance_masks: a numpy array of shape [N, image_height, image_width] with
        values ranging between 0 and 1, can be None.
      instance_boundaries: a numpy array of shape [N, image_height, image_width]
        with values ranging between 0 and 1, can be None.
      keypoints: a numpy array of shape [N, num_keypoints, 2], can
        be None
      use_normalized_coordinates: whether boxes is to be interpreted as
        normalized coordinates or not.
      max_boxes_to_draw: maximum number of boxes to visualize.  If None, draw
        all boxes.
      min_score_thresh: minimum score threshold for a box to be visualized
      agnostic_mode: boolean (default: False) controlling whether to evaluate in
        class-agnostic mode or not.  This mode will display scores but ignore
        classes.
      line_thickness: integer (default: 4) controlling line width of the boxes.
      groundtruth_box_visualization_color: box color for visualizing groundtruth
        boxes
      skip_scores: whether to skip score when drawing a single detection
      skip_labels: whether to skip label when drawing a single detection

    Returns:
      uint8 numpy array with shape (img_height, img_width, 3) with overlaid boxes.
    """
    # Create a display string (and color) for every box location, group any boxes
    # that correspond to the same location.
    box_to_display_str_map = collections.defaultdict(list)
    box_to_color_map = collections.defaultdict(str)
    box_to_instance_masks_map = {}
    box_to_instance_boundaries_map = {}
    box_to_keypoints_map = collections.defaultdict(list)
    if not max_boxes_to_draw:
        max_boxes_to_draw = boxes.shape[0]
    for i in range(min(max_boxes_to_draw, boxes.shape[0])):
        if scores is None or scores[i] >= min_score_thresh:
            box = tuple(boxes[i].tolist())
            if instance_masks is not None:
                box_to_instance_masks_map[box] = instance_masks[i]
            if instance_boundaries is not None:
                box_to_instance_boundaries_map[box] = instance_boundaries[i]
            if keypoints is not None:
                box_to_keypoints_map[box].extend(keypoints[i])
            if False and scores is None:  ## do not go to this branch
                box_to_color_map[box] = groundtruth_box_visualization_color
            else:
                display_str = ''
                if not skip_labels:
                    if not agnostic_mode:
                        if classes[i] in category_index.keys():
                            class_name = category_index[classes[i]]['name']
                        else:
                            class_name = 'N/A'
                        display_str = str(class_name)
                if not skip_scores:
                    if not display_str:
                        display_str = '{}%'.format(int(100 * scores[i]))
                    else:
                        display_str = '{}: {}%'.format(display_str, int(100 * scores[i]))
                box_to_display_str_map[box].append(display_str)
                if agnostic_mode:
                    box_to_color_map[box] = 'DarkOrange'
                else:
                    box_to_color_map[box] = STANDARD_COLORS[
                        classes[i] % len(STANDARD_COLORS)]

    # Draw all boxes onto image.
    for box, color in box_to_color_map.items():
        ymin, xmin, ymax, xmax = box
        if instance_masks is not None:
            draw_mask_on_image_array(
                image,
                box_to_instance_masks_map[box],
                color=color
            )
        if instance_boundaries is not None:
            draw_mask_on_image_array(
                image,
                box_to_instance_boundaries_map[box],
                color='red',
                alpha=1.0
            )
        draw_bounding_box_on_image_array(
            image,
            ymin,
            xmin,
            ymax,
            xmax,
            color=color,
            thickness=line_thickness,
            display_str_list=box_to_display_str_map[box],
            use_normalized_coordinates=use_normalized_coordinates)
        if keypoints is not None:
            draw_keypoints_on_image_array(
                image,
                box_to_keypoints_map[box],
                color=color,
                radius=line_thickness / 2,
                use_normalized_coordinates=use_normalized_coordinates)

    return image


def visualize_polygons_and_labels_on_image_array(
    image,
    polygons,
    classes,
    scores,
    category_index,
    instance_masks=None,
    instance_boundaries=None,
    keypoints=None,
    use_normalized_coordinates=False,
    max_polygons_to_draw=None,
    min_score_thresh=0.,
    agnostic_mode=False,
    line_thickness=4,
    groundtruth_box_visualization_color='black',
    skip_scores=False,
    skip_labels=False):
    """Overlay labeled boxes on an image with formatted scores and label names.

    This function groups boxes that correspond to the same location
    and creates a display string for each detection and overlays these
    on the image. Note that this function modifies the image in place, and returns
    that same image.

    Args:
      image: uint8 numpy array with shape (img_height, img_width, 3)
      boxes: a numpy array of shape [N, 4]
      classes: a numpy array of shape [N]. Note that class indices are 1-based,
        and match the keys in the label map.
      scores: a numpy array of shape [N] or None.  If scores=None, then
        this function assumes that the boxes to be plotted are groundtruth
        boxes and plot all boxes as black with no classes or scores.
      category_index: a dict containing category dictionaries (each holding
        category index `id` and category name `name`) keyed by category indices.
      instance_masks: a numpy array of shape [N, image_height, image_width] with
        values ranging between 0 and 1, can be None.
      instance_boundaries: a numpy array of shape [N, image_height, image_width]
        with values ranging between 0 and 1, can be None.
      keypoints: a numpy array of shape [N, num_keypoints, 2], can
        be None
      use_normalized_coordinates: whether boxes is to be interpreted as
        normalized coordinates or not.
      max_boxes_to_draw: maximum number of boxes to visualize.  If None, draw
        all boxes.
      min_score_thresh: minimum score threshold for a box to be visualized
      agnostic_mode: boolean (default: False) controlling whether to evaluate in
        class-agnostic mode or not.  This mode will display scores but ignore
        classes.
      line_thickness: integer (default: 4) controlling line width of the boxes.
      groundtruth_box_visualization_color: box color for visualizing groundtruth
        boxes
      skip_scores: whether to skip score when drawing a single detection
      skip_labels: whether to skip label when drawing a single detection

    Returns:
      uint8 numpy array with shape (img_height, img_width, 3) with overlaid boxes.
    """
    # Create a display string (and color) for every box location, group any boxes
    # that correspond to the same location.
    idx_to_display_str_map = collections.defaultdict(list)
    idx_to_color_map = collections.defaultdict(str)
    idx_to_instance_masks_map = {}
    idx_to_instance_boundaries_map = {}
    idx_to_keypoints_map = collections.defaultdict(list)
    if not max_polygons_to_draw:
        max_polygons_to_draw = len(polygons)
    for i in range(min(max_polygons_to_draw, len(polygons))):
        if scores is None or scores[i] >= min_score_thresh:
            if instance_masks is not None:
                idx_to_instance_masks_map[i] = instance_masks[i]
            if instance_boundaries is not None:
                idx_to_instance_boundaries_map[i] = instance_boundaries[i]
            if keypoints is not None:
                idx_to_keypoints_map[i].extend(keypoints[i])
            if False and scores is None:  ## do not go to this branch
                idx_to_color_map[i] = groundtruth_box_visualization_color
            else:
                display_str = ''
                if not skip_labels:
                    if not agnostic_mode:
                        if classes[i] in category_index.keys():
                            class_name = category_index[classes[i]]['name']
                        else:
                            class_name = 'N/A'
                        display_str = str(class_name)
                if not skip_scores:
                    if not display_str:
                        display_str = '{}%'.format(int(100 * scores[i]))
                    else:
                        display_str = '{}: {}%'.format(display_str, int(100 * scores[i]))
                idx_to_display_str_map[i].append(display_str)
                if agnostic_mode:
                    idx_to_color_map[i] = 'DarkOrange'
                else:
                    idx_to_color_map[i] = STANDARD_COLORS[
                        classes[i] % len(STANDARD_COLORS)]

    # Draw all polygons onto image.
    for idx, color in idx_to_color_map.items():
        polygon = polygons[idx]
        if instance_masks is not None:
            draw_mask_on_image_array(
                image,
                idx_to_instance_masks_map[idx],
                color=color
            )
        if instance_boundaries is not None:
            draw_mask_on_image_array(
                image,
                idx_to_instance_boundaries_map[idx],
                color='red',
                alpha=1.0
            )
        draw_bounding_plygon_on_image_array(
            image,
            polygon,
            color=color,
            thickness=line_thickness,
            display_str_list=idx_to_display_str_map[idx],
            use_normalized_coordinates=use_normalized_coordinates)
        if keypoints is not None:
            draw_keypoints_on_image_array(
                image,
                idx_to_keypoints_map[idx],
                color=color,
                radius=line_thickness / 2,
                use_normalized_coordinates=use_normalized_coordinates)

    return image


def draw_bounding_plygons_on_image_array(image,
                                         polygon_list,
                                         color_or_color_list='red',
                                         thickness=4,
                                         display_str_list_list=(),
                                         use_normalized_coordinates=True):
    """Adds a bounding box to an image (numpy array).

    Bounding box coordinates can be specified in either absolute (pixel) or
    normalized coordinates by setting the use_normalized_coordinates argument.

    Args:
      image: a numpy array with shape [height, width, 3].
      polygon_list: a list of numpy array with shape [M, 2].
      color_or_color_list: color to draw bounding box. Default is red.
      thickness: line thickness. Default value is 4.
      display_str_list_list: list of strings to display in box
                        (each to be shown on its own line).
      use_normalized_coordinates: If True (default), treat coordinates
        ymin, xmin, ymax, xmax as relative to the image.  Otherwise treat
        coordinates as absolute.
    """
    if isinstance(color_or_color_list, str):
        color_or_color_list = [color_or_color_list] * len(polygon_list)
    assert len(polygon_list) == len(color_or_color_list) == len(display_str_list_list), "args must be same length"
    for plg, clr, dstr in zip(polygon_list, color_or_color_list, display_str_list_list):
        draw_bounding_plygon_on_image_array(image, plg, clr, thickness, dstr,
                                            use_normalized_coordinates)


def draw_bounding_plygon_on_image_array(image,
                                        polygon,
                                        color='red',
                                        thickness=4,
                                        display_str_list=(),
                                        use_normalized_coordinates=True):
    """Adds a bounding box to an image (numpy array).

    Bounding box coordinates can be specified in either absolute (pixel) or
    normalized coordinates by setting the use_normalized_coordinates argument.

    Args:
      image: a numpy array with shape [height, width, 3].
      polygon: a numpy array with shape [M, 2].
      color: color to draw bounding box. Default is red.
      thickness: line thickness. Default value is 4.
      display_str_list: list of strings to display in box
                        (each to be shown on its own line).
      use_normalized_coordinates: If True (default), treat coordinates
        ymin, xmin, ymax, xmax as relative to the image.  Otherwise treat
        coordinates as absolute.
    """
    image_pil = Image.fromarray(np.uint8(image)).convert('RGB')
    draw_bounding_polygon_on_image(image_pil, polygon, color,
                                   thickness, display_str_list,
                                   use_normalized_coordinates)
    np.copyto(image, np.array(image_pil))


def draw_bounding_polygon_on_image(image,
                                   polygon,
                                   color='red',
                                   thickness=4,
                                   display_str_list=(),
                                   use_normalized_coordinates=True):
    """Adds a bounding box to an image.

    Bounding box coordinates can be specified in either absolute (pixel) or
    normalized coordinates by setting the use_normalized_coordinates argument.

    Each string in display_str_list is displayed on a separate line above the
    bounding box in black text on a rectangle filled with the input 'color'.
    If the top of the bounding box extends to the edge of the image, the strings
    are displayed below the bounding box.

    Args:
      image: a PIL.Image object.
      polygon: a numpy array with shape [M, 2], [[x0, y0], [x1, y1] ...].
      color: color to draw bounding box. Default is red.
      thickness: line thickness. Default value is 4.
      display_str_list: list of strings to display in box
                        (each to be shown on its own line).
      use_normalized_coordinates: If True (default), treat coordinates
        ymin, xmin, ymax, xmax as relative to the image.  Otherwise treat
        coordinates as absolute.
    """
    draw = ImageDraw.Draw(image)
    im_width, im_height = image.size
    if use_normalized_coordinates:
        lines = [(x * im_width, y * im_height) for x, y in polygon]
    else:
        lines = [(x, y) for x, y in polygon]
    lines.append(lines[0])
    draw.line(lines, width=thickness, fill=color)

    font = get_default_font(24)

    # If the total height of the display strings added to the top of the bounding
    # box exceeds the top of the image, stack the strings below the bounding box
    # instead of above.
    display_str_heights = [font.getsize(ds)[1] for ds in display_str_list]
    # Each display_str has a top and bottom margin of 0.05x.
    total_display_str_height = (1 + 2 * 0.05) * sum(display_str_heights)

    left, top = np.min(lines, axis=0)
    right, bottom = np.max(lines, axis=0)
    if top > total_display_str_height:
        text_bottom = top
    else:
        text_bottom = bottom + total_display_str_height
    # Reverse list and print from bottom to top.
    for display_str in display_str_list[::-1]:
        text_width, text_height = font.getsize(display_str)
        margin = np.ceil(0.05 * text_height)
        draw.rectangle(
            [(left, text_bottom - text_height - 2 * margin), (left + text_width,
                                                              text_bottom)],
            fill=color)
        draw.text(
            (left + margin, text_bottom - text_height - margin),
            display_str,
            fill='black',
            font=font)
        text_bottom -= text_height - 2 * margin


# draw table on image
txt_cel = lambda s, c: {"type": "text", "name": s, "color": c}
rct_cel = lambda s, c: {"type": "rect", "size": s, "color": c}
plc_cel = lambda: {"type": "placeholder"}


def cal_cell_cord(table_sizes, margin, cell_margin_ratio):
    def cal_row_or_col_cord(lengths, mrg, cel_mrg_r):
        lengths_p = lengths * cel_mrg_r[0]
        lengths_s = lengths * cel_mrg_r[1]
        lengths_p[0] += mrg[0]
        lengths_s[-1] += mrg[1]

        cd = np.cumsum(lengths + lengths_p + lengths_s)
        l = cd[-1]
        cd[1:] = cd[:-1]
        cd[0] = 0

        xs = cd + lengths_p
        xd = xs + lengths

        return xs, l

    col_widths = np.max(table_sizes[:, :, 0], axis=0)
    row_heights = np.max(table_sizes[:, :, 1], axis=1)

    xs, w = cal_row_or_col_cord(col_widths, (margin[0], margin[2]),
                                (cell_margin_ratio[0], cell_margin_ratio[2]))
    ys, h = cal_row_or_col_cord(row_heights, (margin[1], margin[3]),
                                (cell_margin_ratio[1], cell_margin_ratio[3]))
    xs, ys = np.meshgrid(xs, ys)
    xys = np.stack([xs, ys], axis=-1)
    xys_ed = xys + table_sizes
    cord = np.concatenate([xys, xys_ed], axis=-1)

    return cord, (w, h)


def draw_table(table_data, image=None, image_width=None, image_height=None,
               margin=None, cell_margin_ratio=None):
    """
    Draw table base on
    """
    try:
        font_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'DFPKingGothicGB-Light-2.ttf')
        font = ImageFont.truetype(font_path, 24)
    except IOError:
        font = ImageFont.load_default()

    if margin is None:
        margin = (16, 16, 16, 16)
    if cell_margin_ratio is None:
        cell_margin_ratio = (0.1, 0.2, 0.1, 0.2)

    # get image size
    table_sizes = []
    for row in table_data:
        row_sizes = []
        for cel in row:
            if cel['type'] == 'text':
                w, h = font.getsize(cel['name'])
            elif cel['type'] == "rect":
                w, h = cel['size']
            elif cel['type'] == "placeholder":
                w, h = 10, 10
            else:
                raise ValueError('cell type must be "text" or "rect", but get "%s"' % cel['type'])
            row_sizes.append([w, h])
        table_sizes.append(row_sizes)
    table_sizes = np.asarray(table_sizes)

    table_cord, size = cal_cell_cord(table_sizes, margin, cell_margin_ratio)
    size = tuple(int(round(1 + x)) if y is None else y for x, y in zip(size, [image_width, image_height]))

    if image is None:
        image = Image.new('RGB', size, (255, 255, 255))
    draw = ImageDraw.Draw(image)
    for row, r_pos in zip(table_data, table_cord.tolist()):
        for cel, pos in zip(row, r_pos):
            if cel['type'] == 'text':
                draw.text(pos[:2], cel['name'], fill=cel['color'], font=font)
            elif cel['type'] == "rect":
                draw.rectangle([tuple(pos[:2]), tuple(pos[2:])], fill=cel['color'])
            elif cel['type'] == "placeholder":
                pass
            else:
                pass
    return image


def visualize_labels_statistic_result(classes, category_index, image_width=None, image_height=None):
    if isinstance(classes, np.ndarray):
        classes = classes.tolist()
    cls_cnt = collections.Counter(classes)
    keys = sorted(cls_cnt.keys())

    tb = []
    hd = [plc_cel(), txt_cel('NAME', 'Black'), txt_cel('COUNT', 'Black')]
    tb.append(hd)
    rec_size = (24, 24)
    for k in keys:
        color = STANDARD_COLORS[k % len(STANDARD_COLORS)]
        if k in category_index.keys():
            class_name = category_index[k]['name']
        else:
            class_name = 'N/A'
        tb.append([rct_cel(rec_size, color), txt_cel(class_name, 'Black'), txt_cel(str(cls_cnt[k]), 'Red')])

    img = draw_table(tb, image_width=image_width, image_height=image_height)
    return img


def visualize_objects_labels_on_image_array(
    image,
    boxes,
    classes,
    scores,
    category_index,
    instance_masks=None,
    instance_boundaries=None,
    keypoints=None,
    use_normalized_coordinates=False,
    line_thickness=4,
    skip_scores=False,
    skip_labels=False,
    layout="vertical"):
    height, width = image.shape[:2]
    visualize_boxes_and_labels_on_image_array(
        image,
        boxes,
        classes,
        scores,
        category_index,
        instance_masks=instance_masks,
        instance_boundaries=instance_boundaries,
        keypoints=keypoints,
        use_normalized_coordinates=use_normalized_coordinates,
        max_boxes_to_draw=None,
        min_score_thresh=0.,
        agnostic_mode=False,
        line_thickness=line_thickness,
        groundtruth_box_visualization_color='black',
        skip_scores=skip_scores,
        skip_labels=skip_labels)
    if layout == "vertical":
        stat_img = visualize_labels_statistic_result(classes, category_index, image_width=width)
        stat_img = np.asarray(stat_img)
        image = np.concatenate([image, stat_img], axis=0)
    elif layout == "horizontal":
        stat_img = visualize_labels_statistic_result(classes, category_index, image_height=height)
        stat_img = np.asarray(stat_img)
        image = np.concatenate([image, stat_img], axis=1)
    else:
        raise ValueError("layout should be 'vertical' or 'horizontal', but get '%s'" % layout)
    return image


def draw_match_result_of_boxes_on_image_array(image, boxes1, boxes2, match_result, thickness=4):
    """Draw matching result of two box list to an image array.

    Bounding box coordinates must be specified in absolute (pixel) coordinates.

    Match result specifies the match pair between boxes1 and boxes2 list and their match result,
    when the paired boxes have different labels, the match result is False, else is True.

    Args:
      image: a numpy array.
      boxes1: a numpy array with shape [M, 4].
      boxes2: a numpy array with shape [M, 4].
      match_result: a list of (index1, index2, Match or not)
      thickness: line thickness. Default value is 4.
    """
    radius = 2 * thickness

    cent_yx1 = np.mean(boxes1.reshape([-1, 2, 2]), axis=1)
    cent_yx2 = np.mean(boxes2.reshape([-1, 2, 2]), axis=1)

    image_pil = Image.fromarray(image)
    draw = ImageDraw.Draw(image_pil)
    for idx1, idx2, res in match_result:
        # color = "Green" if res else "Red"
        clmp = {"correct": "Blue", "error": "Red", "igr_box": "Grey", "igr_region": "Grey"}
        color = clmp[res]
        line = [tuple(cent_yx1[idx1, ::-1]), tuple(cent_yx2[idx2, ::-1])]
        # draw.line(line, fill=color, width=thickness)
        draw.ellipse([(line[0][0] - radius, line[0][1] - radius),
                      (line[0][0] + radius, line[0][1] + radius)],
                     outline=color, fill=color)
        draw.ellipse([(line[1][0] - radius, line[1][1] - radius),
                      (line[1][0] + radius, line[1][1] + radius)],
                     outline=color, fill=color)

    match1 = set([x[0] for x in match_result])
    miss_match1 = [i for i in range(len(boxes1)) if i not in match1]
    for idx1 in miss_match1:
        ctxy = cent_yx1[idx1, ::-1]
        draw.ellipse([(ctxy[0] - radius, ctxy[1] - radius),
                      (ctxy[0] + radius, ctxy[1] + radius)],
                     outline="Yellow", fill="Yellow")

    match2 = set([x[1] for x in match_result])
    miss_match2 = [i for i in range(len(boxes2)) if i not in match2]
    for idx2 in miss_match2:
        ctxy = cent_yx2[idx2, ::-1]
        draw.ellipse([(ctxy[0] - radius, ctxy[1] - radius),
                      (ctxy[0] + radius, ctxy[1] + radius)],
                     outline="Yellow", fill="Yellow")

    np.copyto(image, np.array(image_pil))
    return image


def visualize_groudtruth_and_predict_on_image_array(
    image,
    groudtruth_boxes,
    groudtruth_classes,
    predict_boxes,
    predict_classes,
    predict_scores,
    category_index,
    match_result,
    use_normalized_coordinates=False):
    """

    Args:
      image: a numpy array.
      boxes1: a numpy array with shape [M, 4].
      boxes2: a numpy array with shape [M, 4].
      match_result: a list of (index1, index2, Match or not)
      thickness: line thickness. Default value is 4.
    """
    bar_length = 10

    height, width = image.shape[:2]
    if use_normalized_coordinates:
        sizes = np.asarray([height, width, height, width])
        groudtruth_boxes = groudtruth_boxes * sizes
        predict_boxes = predict_boxes * sizes

    gt_img = visualize_objects_labels_on_image_array(np.copy(image), groudtruth_boxes,
                                                     groudtruth_classes, None,
                                                     category_index,
                                                     skip_scores=True,
                                                     skip_labels=True)
    pd_img = visualize_objects_labels_on_image_array(np.copy(image), predict_boxes,
                                                     predict_classes, predict_scores,
                                                     category_index,
                                                     skip_scores=True if predict_scores is None else False,
                                                     skip_labels=True)
    gt_shape = gt_img.shape
    pd_shape = pd_img.shape

    # shape = [gt_shape[0] + pd_shape[0] + bar_length, max(gt_shape[1], pd_shape[1]), gt_shape[2]]
    # res_img = np.ones(shape, dtype=gt_img.dtype) * 255
    # res_img[:gt_shape[0], :gt_shape[1], :] = gt_img
    # res_img[gt_shape[0]:gt_shape[0]+bar_length, :, :] = 0
    # res_img[gt_shape[0] + bar_length:gt_shape[0]+ bar_length+pd_shape[0], :pd_shape[1], :] = pd_img
    # # draw match result
    # predict_boxes[:, [0, 2]] += gt_shape[0] + bar_length

    shape = [max(gt_shape[0], pd_shape[0]), gt_shape[1] + pd_shape[1] + bar_length, gt_shape[2]]
    res_img = np.ones(shape, dtype=gt_img.dtype) * 255
    res_img[:gt_shape[0], :gt_shape[1], :] = gt_img
    res_img[:, gt_shape[1]:gt_shape[1] + bar_length, :] = 0
    res_img[:pd_shape[0], gt_shape[1] + bar_length:gt_shape[1] + bar_length + pd_shape[1], :] = pd_img
    # draw match result
    predict_boxes[:, [1, 3]] += gt_shape[1] + bar_length

    draw_match_result_of_boxes_on_image_array(res_img, groudtruth_boxes, predict_boxes, match_result)

    return res_img
